from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM,TrainingArguments,Trainer,DataCollatorForSeq2Seq
from modelscope import AutoTokenizer
from datasets import Dataset
import torch
import swanlab
import os
import pandas as pd

# os.environ["CUDA_VISIBLE_DEVICES"] = ""  
os.environ["SWANLAB_PROJECT"] = 'qwen3-sft-medical'

PROMPT = '你是一个医学专家，你需要根据用户的问题，一步一步思考，给出带有思考的回答。'
MAX_LENGTH = 2048

swanlab.config.update(
    {
        'model':'Qwen/Qwen3-0.6B-lora',
        'prompt': PROMPT,
        'data_max_length':MAX_LENGTH,
    }
)

# 加载模型和tokenizer
model_folder = r"D:\models\qwen3-0.6b"
model = AutoModelForCausalLM.from_pretrained(model_folder,device_map='auto',torch_dtype='auto')
tokenizer = AutoTokenizer.from_pretrained(model_folder, use_fast=False, trust_remote_code=True)
model.enable_input_require_grads() 
# 明确告知基础模型为其输入启用梯度计算，即使其参数已冻结。
# 配置 LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM, # 这是因果语言模型任务
    inference_mode=False,
    r=8, # LoRA 秩，通常是 8, 16, 32, 64。越大参数越多，性能越好，但内存消耗也越大
    lora_alpha=16, # lora_alpha 配合 r，通常设为 2*r 或更大
    lora_dropout=0.05, # 防止过拟合
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], # Qwen3-1.7B 中常见的可适配模块
    # Qwen3DecoderLayer.self_attn.q_proj, k_proj, v_proj, o_proj
    # Qwen3MLP.gate_proj, up_proj, down_proj
)

# 将模型包装成 PEFT 模型
model = get_peft_model(model, lora_config)
model.print_trainable_parameters() # 打印可训练参数量，你会发现它非常小

# 数据加载
def process_func(example):
    """数据集预处理"""
    input_ids, attention_mask, labels = [], [], []
    instruction = tokenizer(
        f"<|im_start|>system\n{PROMPT}<|im_end|>\n<|im_start|>user\n{example['input']}<|im_end|>\n<|im_start|>assistant\n",
        add_special_tokens=False,
    )
    # print(f'instruction: {instruction[0]} \n')
    response = tokenizer(f"{example['output']}", add_special_tokens=False)
    input_ids = instruction["input_ids"] + response["input_ids"] #+ [tokenizer.pad_token_id]
    attention_mask = (
        instruction["attention_mask"] + response["attention_mask"] # +[1]
    )
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] #+ [tokenizer.pad_token_id]
    if len(input_ids) > MAX_LENGTH:  # 做一个截断
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}   

train_jsonl_new_path = "train_format.jsonl"
test_jsonl_new_path = "val_format.jsonl"

# 得到训练集
train_df = pd.read_json(train_jsonl_new_path, lines=True)
train_ds = Dataset.from_pandas(train_df)
train_dataset = train_ds.map(process_func, remove_columns=train_ds.column_names)

# 得到验证集
eval_df = pd.read_json(test_jsonl_new_path, lines=True)
eval_ds = Dataset.from_pandas(eval_df)
eval_dataset = eval_ds.map(process_func, remove_columns=eval_ds.column_names)

# Trainer 实例化 训练
args = TrainingArguments(
    output_dir=r"C:\Users\COLORFUL\Desktop\AI_NLP\大模型微调\lora-qwen\model",
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=1,
    eval_strategy="steps",
    eval_steps=100,
    logging_steps=10,
    num_train_epochs=2,
    save_steps=200,
    learning_rate=5e-4,
    save_on_each_node=True,
    gradient_checkpointing=True,
    report_to="swanlab",
    run_name="lora-qwen3-0.6B-lora",
)

trainer = Trainer(
    model = model,
    args = args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
)

trainer.train()

def predict(messages, model, tokenizer):
    device = "cpu"
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
 
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    generated_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=MAX_LENGTH,
        # temperature=0.7,  # 添加温度参数使生成更自然
        # top_p=0.85,       # 添加top-p采样
        # do_sample=True,   # 启用采样
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True,clean_up_tokenization_spaces=True)[0]

    return response

# 测试集的前3条，主观看模型效果
test_df = pd.read_json(test_jsonl_new_path, lines=True)[:3]
test_text_list = []

model.eval()

for index, row in test_df.iterrows():
    instruction = row['instruction']
    input_value = row['input']

    messages = [
        {"role": "system", "content": f"{instruction}"},
        {"role": "user", "content": f"{input_value}"}
    ]

    response = predict(messages, model, tokenizer)

    response_text = f"""
    Question: {input_value}

    LLM:{response}
    """
    
    test_text_list.append(swanlab.Text(response_text))
    print(response_text)

swanlab.log({"Prediction": test_text_list})

swanlab.finish()